#include "pch.h"
#include "AsyncLibImpl.h"
#include "TaskThread\MainTaskThread.h"
#include "TaskThread\ChildTaskThread.h"
#include "Threadpool\ThreadpoolWork.h"
#include "Threadpool\ThreadpoolTimer.h"
#include "Threadpool\ThreadpoolWait.h"
#include "Threadpool\TpTimerManager.h"
#include "Threadpool\TpFileManager.h"

namespace AsyncLib
{

static_assert(ATID_MAIN == 1);

thread_local intptr_t g_nAtid;

CAsyncLib::CAsyncLib() : m_status(Exited), m_tpCallbackEnv{}, m_pTpPool(NULL), m_pTpCleanupGroup(NULL)
{
}

CAsyncLib::~CAsyncLib()
{
	assert(m_status == Exited);
}

void CAsyncLib::Init(ATID atidMax)
{
	std::unique_lock lock(m_mutex);
	assert(m_status == Exited);

	if (m_status != Exited)
		return;

	m_status = Running;
	SetInTaskThread(ATID_MAIN);
	InitTaskThreads(atidMax);
	InitThreadpools();
}

void CAsyncLib::Exit(bool bWaitTasks)
{
	if (bWaitTasks)
		WaitForAllTasks();

	std::unique_lock lock(m_mutex);
	assert(m_status == Running);

	if (m_status != Running)
		return;

	m_status = Exiting;
	lock.unlock();
	ExitTaskThreads();
	ExitThreadpools();
	ProcessCallbacks();
}

void CAsyncLib::WaitForAllTasks()
{
	std::shared_lock lock(m_mutex);

	if (!IsRunning())
		return;

	size_t nRefCount = m_vecTaskThreads.size();
	std::condition_variable_any condVar;

	auto fnCallback = [this, &nRefCount, &condVar]
	{
		std::unique_lock lock(m_mutex);

		if (--nRefCount == 0)
			condVar.notify_one();
	};

	// Post the cut-off task to all members
	for (size_t i = 1; i < m_vecTaskThreads.size(); i++)
	{
		m_vecTaskThreads[i]->PostTask(fnCallback);
	}

	m_pThreadpoolWork->PostTask(fnCallback);

	// Wait for all cut-off tasks to execute
	while (nRefCount)
	{
		condVar.wait(lock);
	}
}

void CAsyncLib::InitTaskThreads(ATID atidMax)
{
	size_t nIndex = atidMax - ATID_MAIN;
	assert(nIndex < 64);
	m_vecTaskThreads.reserve(nIndex + 1);

	auto pMainTaskThread = std::make_unique<CMainTaskThread>();
	pMainTaskThread->CreateWnd();
	m_vecTaskThreads.emplace_back(std::move(pMainTaskThread));

	for (size_t i = 1; i <= nIndex; i++)
	{
		m_vecTaskThreads.emplace_back(std::make_unique<CChildTaskThread>(ATID(i + ATID_MAIN)));
	}
}

void CAsyncLib::ExitTaskThreads()
{
	for (auto& pTaskThread : m_vecTaskThreads)
	{
		pTaskThread->Stop();
	}

	for (auto& pTaskThread : m_vecTaskThreads)
	{
		pTaskThread->Exit();
	}

	auto vecTaskThreads = std::move(m_vecTaskThreads);
}

void CAsyncLib::InitThreadpools()
{
	InitializeThreadpoolEnvironment(&m_tpCallbackEnv);

	SYSTEM_INFO sysInfo;
	GetSystemInfo(&sysInfo);
	UINT nMaxThreads = sysInfo.dwNumberOfProcessors, nMinThreads = nMaxThreads / 2;

	m_pTpPool = CreateThreadpool(NULL);
	SetThreadpoolThreadMinimum(m_pTpPool, nMinThreads);
	SetThreadpoolThreadMaximum(m_pTpPool, nMaxThreads);
	SetThreadpoolCallbackPool(&m_tpCallbackEnv, m_pTpPool);

	m_pTpCleanupGroup = CreateThreadpoolCleanupGroup();
	SetThreadpoolCallbackCleanupGroup(&m_tpCallbackEnv, m_pTpCleanupGroup, NULL);

	m_pThreadpoolWork = std::make_unique<CThreadpoolWork>(nMinThreads, nMaxThreads);
	m_pThreadpoolWork->Init();
	m_pThreadpoolTimer = std::make_unique<CThreadpoolTimer>();
	m_pThreadpoolTimer->Init();
	m_pThreadpoolWait = std::make_unique<CThreadpoolWait>();
	m_pThreadpoolWait->Init();
	m_pTpTimerManager = std::make_shared<CTpTimerManager>();
	m_pTpTimerManager->Init();
	m_pTpFileManager = std::make_shared<CTpFileManager>();
	m_pTpFileManager->Init();
}

void CAsyncLib::ExitThreadpools()
{
	// All members of m_pTpCleanupGroup must be locked and IsRunning() when used
	m_pThreadpoolWork->Stop();
	m_pThreadpoolTimer->Stop();
	m_pThreadpoolWait->Stop();
	m_pTpTimerManager->Stop();
	m_pTpFileManager->Stop();
	CloseThreadpoolCleanupGroupMembers(m_pTpCleanupGroup, TRUE, NULL);

	m_pThreadpoolWork->Exit();
	m_pThreadpoolWork = nullptr;
	m_pThreadpoolTimer->Exit();
	m_pThreadpoolTimer = nullptr;
	m_pThreadpoolWait->Exit();
	m_pThreadpoolWait = nullptr;
	m_pTpTimerManager->Exit();
	m_pTpTimerManager = nullptr;
	m_pTpFileManager->Exit();
	m_pTpFileManager = nullptr;

	CloseThreadpoolCleanupGroup(m_pTpCleanupGroup);
	CloseThreadpool(m_pTpPool);
	DestroyThreadpoolEnvironment(&m_tpCallbackEnv);
	m_pTpCleanupGroup = NULL;
	m_pTpPool = NULL;
}

void CAsyncLib::PostTask(ATID atid, std::function<void()>&& fnTask)
{
	assert(atid);
	size_t nIndex = atid - ATID_MAIN;
	std::shared_lock lock(m_mutex);

	if (IsRunning() && fnTask)
	{
		assert(nIndex < m_vecTaskThreads.size() || atid == ATID_TP);

		if (nIndex < m_vecTaskThreads.size())
			m_vecTaskThreads[nIndex]->PostTask(std::move(fnTask));
		else
			m_pThreadpoolWork->PostTask(std::move(fnTask));
	}
}

void CAsyncLib::PostDelayedTask(UINT nMilliseconds, std::function<void()>&& fnTask)
{
	std::shared_lock lock(m_mutex);

	if (IsRunning() && fnTask)
		m_pThreadpoolTimer->AddTask(nMilliseconds, std::move(fnTask));
}

void CAsyncLib::WaitForAsync(HANDLE handle, UINT nMilliseconds, std::function<void(DWORD)>&& fnCallback)
{
	std::shared_lock lock(m_mutex);

	if (IsRunning() && handle && fnCallback)
		m_pThreadpoolWait->AddTask(handle, nMilliseconds, std::move(fnCallback));
}

bool CAsyncLib::IsInTaskThread(ATID atid) const
{
	assert(atid);
	return g_nAtid == atid || m_status == Exited;
}

bool CAsyncLib::IsInThreadPool() const
{
	return IsInTaskThread(ATID_TP);
}

void CAsyncLib::SetInTaskThread(ATID atid)
{
	assert(atid);
	g_nAtid = atid;
}

void CAsyncLib::SetInThreadPool()
{
	SetInTaskThread(ATID_TP);
}

ATID CAsyncLib::InnerGetAtid() const
{
	return ATID(g_nAtid);
}

void CAsyncLib::InnerPostTask(ATID atid, std::function<void()>&& fnTask)
{
	assert(atid);
	size_t nIndex = atid - ATID_MAIN;
	std::unique_lock lock(m_mutex);

	// Cannot PostTask while exiting, and cannot call immediately because of contention
	switch (m_status)
	{
	case Running:
		if (nIndex < m_vecTaskThreads.size())
			m_vecTaskThreads[nIndex]->PostTask(std::move(fnTask));
		else
			m_pThreadpoolWork->PostTask(std::move(fnTask));
		break;

	case Exiting:
		m_deqTasks.emplace_back(std::move(fnTask));
		break;
	}
}

void CAsyncLib::ProcessCallbacks()
{
	// All threads have exited, we call all callbacks of CaptureThread in order
	std::unique_lock lock(m_mutex);

	while (m_deqTasks.size())
	{
		auto fnTask = std::move(m_deqTasks.front());
		m_deqTasks.pop_front();
		lock.unlock();

		fnTask();
		fnTask = nullptr;

		lock.lock();
	}

	auto deqTasks = std::move(m_deqTasks);
	m_status = Exited;
}

std::shared_ptr<ITimerTask> CAsyncLib::CreateTimerTask(bool bUITimer)
{
	std::shared_lock lock(m_mutex);

	if (!IsRunning())
		return nullptr;
	else if (bUITimer)
		return ((CMainTaskThread&)*m_vecTaskThreads[0]).CreateTimer();
	else
		return m_pTpTimerManager->CreateTimer(m_pTpTimerManager);
}

std::shared_ptr<IFileTask> CAsyncLib::CreateFileTask(HANDLE hFile)
{
	std::shared_lock lock(m_mutex);

	if (!IsRunning())
		return nullptr;
	else
		return m_pTpFileManager->CreateTask(m_pTpFileManager, hFile);
}

} // namespace AsyncLib
